from typing import Dict, List, Optional, Any
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import uvicorn
import yaml
import logging
from pathlib import Path

from core.runtime_info import RuntimeInfo
from core.scheduler_base import SchedulerBase
from core.scheduler_pss import PSSScheduler
from core.task_assigner import TaskAssigner
from core.garbage_collector import GarbageCollector
from .runtime_monitor import RuntimeMonitor

# API Models
class WorkerRegistration(BaseModel):
    worker_id: str
    capabilities: Dict[str, Any]

class TaskUpdate(BaseModel):
    worker_id: str
    task_id: str
    status: str
    result: Optional[Dict[str, Any]] = None

class MetricsReport(BaseModel):
    worker_id: str
    metrics: Dict[str, Any]

class SyncRequest(BaseModel):
    worker_id: str
    sync_point: int
    worker_group: List[str]

class MasterNode:
    """Master node implementation for the distributed system.
    
    This class coordinates all worker nodes, manages task distribution,
    and provides the REST API interface for system control.
    """
    
    def __init__(self, config_path: str):
        """Initialize the master node.
        
        Args:
            config_path: Path to configuration file
        """
        # Load configuration
        with open(config_path, 'r') as f:
            self.config = yaml.safe_load(f)
            
        # Setup logging
        self._setup_logging()
        
        # Initialize components
        self.runtime_info = RuntimeInfo()
        self.scheduler = self._create_scheduler()
        self.task_assigner = TaskAssigner(self.runtime_info, self.scheduler)
        self.garbage_collector = GarbageCollector(
            runtime_info=self.runtime_info,
            task_assigner=self.task_assigner,
            config=self.config
        )
        self.monitor = RuntimeMonitor(self.config)
        
        # Create FastAPI application
        self.app = FastAPI(title="DeepSlicing Master")
        self._setup_routes()
    
    def _setup_logging(self) -> None:
        """Setup logging configuration."""
        logging.basicConfig(
            level=self.config['logging']['level'],
            format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
            filename=self.config['logging']['file']
        )
        self.logger = logging.getLogger("MasterNode")
    
    def _create_scheduler(self) -> SchedulerBase:
        """Create the task scheduler based on configuration.
        
        Returns:
            SchedulerBase: Configured scheduler instance
        """
        scheduler_type = self.config['scheduler']['type']
        if scheduler_type == 'PSS':
            return PSSScheduler(self.config)
        else:
            raise ValueError(f"Unsupported scheduler type: {scheduler_type}")
    
    def _setup_routes(self) -> None:
        """Setup FastAPI routes."""
        # Worker registration
        @self.app.post("/register")
        async def register_worker(registration: WorkerRegistration):
            try:
                self.monitor.register_worker(
                    registration.worker_id,
                    registration.capabilities
                )
                return {"status": "success"}
            except Exception as e:
                raise HTTPException(status_code=500, detail=str(e))
        
        # Task management
        @self.app.get("/get_task")
        async def get_task(worker_id: str):
            try:
                assignments = self.task_assigner.assign_pending_tasks()
                if assignments:
                    assignment = next(
                        (a for a in assignments if a['worker_id'] == worker_id),
                        None
                    )
                    if assignment:
                        self.monitor.update_task_status(
                            task_id=assignment['task_id'],
                            status='running',
                            worker_id=worker_id
                        )
                        return assignment
                return None
            except Exception as e:
                raise HTTPException(status_code=500, detail=str(e))
        
        @self.app.post("/update_task")
        async def update_task(update: TaskUpdate):
            try:
                if update.status == 'completed':
                    self.task_assigner.handle_task_completion(
                        update.task_id,
                        update.worker_id,
                        update.result or {}
                    )
                elif update.status == 'failed':
                    self.task_assigner.handle_task_failure(
                        update.task_id,
                        update.worker_id,
                        update.result or {}
                    )
                
                self.monitor.update_task_status(
                    task_id=update.task_id,
                    status=update.status,
                    worker_id=update.worker_id,
                    error=update.result.get('error') if update.result else None
                )
                
                return {"status": "success"}
            except Exception as e:
                raise HTTPException(status_code=500, detail=str(e))
        
        # Metrics reporting
        @self.app.post("/report_metrics")
        async def report_metrics(report: MetricsReport):
            try:
                self.monitor.update_worker_status(
                    report.worker_id,
                    report.metrics
                )
                return {"status": "success"}
            except Exception as e:
                raise HTTPException(status_code=500, detail=str(e))
        
        # Synchronization
        @self.app.post("/sync_ready")
        async def sync_ready(request: SyncRequest):
            try:
                all_ready = self.monitor.worker_ready_for_sync(
                    request.worker_id,
                    request.sync_point
                )
                return {"success": all_ready}
            except Exception as e:
                raise HTTPException(status_code=500, detail=str(e))
        
        # System status
        @self.app.get("/status")
        async def get_status():
            try:
                return self.monitor.get_system_status()
            except Exception as e:
                raise HTTPException(status_code=500, detail=str(e))
    
    async def start(self) -> None:
        """Start the master node."""
        self.logger.info("Starting master node...")
        
        # Initialize scheduler
        if not self.scheduler.initialize():
            raise RuntimeError("Failed to initialize scheduler")
        
        # Start monitoring
        import asyncio
        asyncio.create_task(self.monitor.start_monitoring())
        
        # Start API server
        config = self.config['system']['master']
        server_config = uvicorn.Config(
            self.app,
            host=config['host'],
            port=config['port'],
            log_level="info"
        )
        server = uvicorn.Server(server_config)
        await server.serve()

def main():
    """Main entry point for the master process."""
    import argparse
    
    parser = argparse.ArgumentParser(description='Start the DeepSlicing master node')
    parser.add_argument('--config', type=str, required=True,
                      help='Path to configuration file')
    
    args = parser.parse_args()
    
    # Create and start master node
    master = MasterNode(args.config)
    
    try:
        import asyncio
        asyncio.run(master.start())
    except KeyboardInterrupt:
        print("\nShutting down master node...")
    except Exception as e:
        print(f"Error: {str(e)}")
        raise

if __name__ == '__main__':
    main()